{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook to demonstrate TAO Object Detection workflow\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "![image](https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png)\n",
    "\n",
    "### Sample prediction for an Object Detection model\n",
    "<img align=\"center\" src=\"../example_images/sample_object_detection.jpg\" width=\"960\">\n",
    "\n",
    "### The workflow in a nutshell\n",
    "\n",
    "- Creating a dataset\n",
    "- Upload dataset to the service\n",
    "- Running dataset convert\n",
    "- Getting a PTM from NGC\n",
    "- Model Actions\n",
    "    - Train (Normal/AutoML)\n",
    "    - Evaluate\n",
    "    - Prune, retrain\n",
    "    - Export\n",
    "    - Tao-Deploy\n",
    "    - Inference on TAO\n",
    "    - Inference on TRT\n",
    "    \n",
    "### Table of contents\n",
    "\n",
    "1. [Create datasets ](#head-1)\n",
    "1. [List the created datasets](#head-2)\n",
    "1. [Dataset convert Action](#head-3)\n",
    "1. [Create an experiment](#head-4)\n",
    "1. [List experiments](#head-5)\n",
    "1. [Assign train, eval datasets](#head-6)\n",
    "1. [Assign PTM](#head-7)\n",
    "1. [View hyperparameters that are enabled by default](#head-8)\n",
    "1. [Set AutoML related configurations](#head-9)\n",
    "1. [Actions](#head-10)\n",
    "1. [Train](#head-11)\n",
    "1. [Evaluate](#head-12)\n",
    "1. [Optimize: Apply specs for prune](#head-14)       \n",
    "1. [Optimize: Apply specs for retrain](#head-15)\n",
    "1. [Optimize: Run actions](#head-16)\n",
    "1. [Export](#head-17)\n",
    "1. [TRT Engine generation using TAO-Deploy](#head-19)\n",
    "1. [TAO inference](#head-20)\n",
    "1. [TRT inference](#head-21)\n",
    "\n",
    "### Requirements\n",
    "Please find the server requirements [here](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_setup.html#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import requests\n",
    "import time\n",
    "import subprocess\n",
    "from IPython.display import clear_output\n",
    "import glob"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FIXME\n",
    "\n",
    "1. Assign a model_name in FIXME 1\n",
    "1. Assign a workdir in FIXME 2\n",
    "1. Assign the ip_address and port_number in FIXME 3 ([info](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_rest_api.html))\n",
    "1. Assign the ngc_api_key variable in FIXME 4\n",
    "1. (Optional) Enable AutoML if needed in FIXME 5\n",
    "1. (Optional) Choose between bayesian and hyperband automl_algorithm in FIXME 6 (If automl was enabled in FIXME5)\n",
    "1. Choose to download jobs or not in FIXME 7\n",
    "1. Choose between default and custom dataset in FIXME 8\n",
    "1. Assign path of DATA_DIR in FIXME 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Define model_name workspaces and other variables\n",
    "# Available models (#FIXME 1):\n",
    "# 1. deformable_detr - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/deformable_detr.html\n",
    "# 2. detectnet_v2 - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/detectnet_v2.html\n",
    "# 3. dino - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/dino.html\n",
    "# 4. dssd - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/ssd.html\n",
    "# 5. efficientdet_tf1 - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/efficientdet_tf1.html\n",
    "# 6. efficientdet_tf2 - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/efficientdet_tf2.html\n",
    "# 7. faster_rcnn - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/fasterrcnn.html\n",
    "# 8. retinanet - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/retinanet.html\n",
    "# 9. ssd - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/ssd.html\n",
    "# 10. yolo_v3 - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/yolo_v3.html\n",
    "# 11. yolo_v4 - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/yolo_v4.html\n",
    "# 12. yolo_v4_tiny - https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/yolo_v4_tiny.html\n",
    "\n",
    "model_name = \"detectnet_v2\" # FIXME1 (Add the model name from the above mentioned list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "workdir = \"workdir_object_detection\" # FIXME2\n",
    "host_url = \"http://<ip_address>:<port_number>\" # FIXME3 example: https://10.137.149.22:32334\n",
    "# In host machine, node ip_address and port number can be obtained as follows,\n",
    "# ip_address: hostname -i\n",
    "# port_number: kubectl get service ingress-nginx-controller -o jsonpath='{.spec.ports[0].nodePort}'\n",
    "ngc_api_key = \"<ngc_api_key>\" # FIXME4 example: (Add NGC API key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "automl_enabled = False # FIXME5 set to True if you want to run automl for the model chosen in the previous cell\n",
    "automl_algorithm = \"bayesian\" # FIXME6 example: bayesian/hyperband\n",
    "# FIXME7 Defaulted to False as downloading jobs from service to your machine takes time\n",
    "# Set to True if you want to download jobs where examples have been provided like for train, export, inference.\n",
    "download_jobs = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Exchange NGC_API_KEY for JWT\n",
    "data = json.dumps({\"ngc_api_key\": ngc_api_key})\n",
    "response = requests.post(f\"{host_url}/api/v1/login\", data=data)\n",
    "user_id = response.json()[\"user_id\"]\n",
    "print(\"User ID\",user_id)\n",
    "token = response.json()[\"token\"]\n",
    "print(\"JWT\",token)\n",
    "\n",
    "# Set base URL\n",
    "base_url = f\"{host_url}/api/v1/users/{user_id}\"\n",
    "print(\"API Calls will be forwarded to\",base_url)\n",
    "\n",
    "headers = {\"Authorization\": f\"Bearer {token}\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Creating workdir\n",
    "if not os.path.isdir(workdir):\n",
    "    os.makedirs(workdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function to split tar files <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tarfile\n",
    "\n",
    "def split_tar_file(input_tar_path, output_dir, max_split_size=0.2*1024*1024*1024):\n",
    "\tos.makedirs(output_dir, exist_ok=True)\n",
    "\t\n",
    "\twith tarfile.open(input_tar_path, 'r') as original_tar:\n",
    "\t\tmembers = original_tar.getmembers()\n",
    "\t\tcurrent_split_size = 0\n",
    "\t\tcurrent_split_number = 0\n",
    "\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\n",
    "\t\twith tarfile.open(current_split_name, 'w') as split_tar:\n",
    "\t\t\tfor member in members:\n",
    "\t\t\t\tif current_split_size + member.size <= max_split_size:\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size\n",
    "\t\t\t\telse:\n",
    "\t\t\t\t\tsplit_tar.close()\n",
    "\t\t\t\t\tcurrent_split_number += 1\n",
    "\t\t\t\t\tcurrent_split_name = os.path.join(output_dir, f'smaller_file_{current_split_number}.tar')\n",
    "\t\t\t\t\tcurrent_split_size = 0\n",
    "\t\t\t\t\tsplit_tar = tarfile.open(current_split_name, 'w')  # Open a new split tar archive\n",
    "\t\t\t\t\tsplit_tar.addfile(member, original_tar.extractfile(member))\n",
    "\t\t\t\t\tcurrent_split_size += member.size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set dataset type, format <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using NVIDIA created `Synthetic Object detection data` based on KITTI dataset format in this notebook. To find more details about kitti format, please visit [here](http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**If using custom dataset; it should follow this dataset structure**\n",
    "```\n",
    "$DATA_DIR/train\n",
    "├── images\n",
    "│   ├── image_name_1.jpg\n",
    "│   ├── image_name_2.jpg\n",
    "|   ├── ...\n",
    "└── labels\n",
    "    ├── image_name_1.txt\n",
    "    ├── image_name_2.txt\n",
    "    ├── ...\n",
    "$DATA_DIR/val\n",
    "├── images\n",
    "│   ├── image_name_5.jpg\n",
    "│   ├── image_name_6.jpg\n",
    "|   ├── ...\n",
    "└── labels\n",
    "    ├── image_name_5.txt\n",
    "    ├── image_name_6.txt\n",
    "    ├── ...\n",
    "```\n",
    "The file name should be same for images and labels folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ds_format = \"kitti\"\n",
    "if model_name in (\"efficientdet_tf1\", \"efficientdet_tf2\", \"deformable_detr\", \"dino\"):\n",
    "    ds_format = \"coco\"\n",
    "ds_type = \"object_detection\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset_to_be_used = \"default\" #FIXME8 #default/custom; default for the dataset used in this tutorial notebook; custom for a different dataset\n",
    "DATA_DIR = model_name # FIXME9\n",
    "os.environ['DATA_DIR']= DATA_DIR\n",
    "!mkdir -p $DATA_DIR\n",
    "job_map = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset download and pre-processing <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if dataset_to_be_used == \"default\":\n",
    "    !python3 -m pip install --upgrade awscli\n",
    "    !aws s3 cp --no-sign-request s3://tao-object-detection-synthetic-dataset/tao_od_synthetic_train.tar.gz $DATA_DIR/\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/tao_od_synthetic_train.tar.gz\"))\n",
    "    !aws s3 cp --no-sign-request s3://tao-object-detection-synthetic-dataset/tao_od_synthetic_val.tar.gz $DATA_DIR/\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/tao_od_synthetic_val.tar.gz\"))\n",
    "\n",
    "    print(\"Untarring file\")\n",
    "    os.makedirs(f\"{DATA_DIR}/train\", exist_ok=True)\n",
    "    !tar -xzf {DATA_DIR}/tao_od_synthetic_train.tar.gz -C {DATA_DIR}/train\n",
    "    os.makedirs(f\"{DATA_DIR}/val\", exist_ok=True)\n",
    "    !tar -xzf {DATA_DIR}/tao_od_synthetic_val.tar.gz -C {DATA_DIR}/val\n",
    "\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/train/images\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/train/labels\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/val/images\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/val/labels\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if ds_format == \"coco\":\n",
    "    !python3 -m pip install ujson opencv-python tqdm\n",
    "    if not os.path.exists(os.path.join(DATA_DIR, \"train\")):\n",
    "        raise Exception(\"Train dataset not present\")\n",
    "    if not os.path.exists(os.path.join(DATA_DIR, \"val\")):\n",
    "        raise Exception(\"Eval dataset not present\")\n",
    "    \n",
    "    #kitti to coco conversion for efficientdet\n",
    "    if model_name == \"efficientdet_tf2\":\n",
    "        label_map_extension = \"yaml\"\n",
    "    else:\n",
    "        label_map_extension = \"txt\"\n",
    "    num_classes = subprocess.getoutput(f'python3 ../dataset_prepare/kitti/kitti_to_coco.py {DATA_DIR}/train/labels {DATA_DIR}/train {label_map_extension}')\n",
    "    subprocess.getoutput(f'python3 ../dataset_prepare/kitti/kitti_to_coco.py {DATA_DIR}/val/labels {DATA_DIR}/val {label_map_extension}')\n",
    "\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/train/images\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/train/annotations.json\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/train/label_map.{label_map_extension}\"))\n",
    "\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/val/images\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/val/annotations.json\"))\n",
    "    assert (os.path.exists(f\"{DATA_DIR}/val/label_map.{label_map_extension}\"))\n",
    "\n",
    "    !tar -C {DATA_DIR}/train -czf {DATA_DIR}/tao_od_synthetic_train.tar.gz images annotations.json label_map.{label_map_extension}\n",
    "    !tar -C {DATA_DIR}/val -czf {DATA_DIR}/tao_od_synthetic_val.tar.gz images annotations.json label_map.{label_map_extension}\n",
    "else:\n",
    "    if dataset_to_be_used == \"custom\":\n",
    "        !tar -C $DATA_DIR/train -czf $DATA_DIR/tao_od_synthetic_train.tar.gz images labels\n",
    "        !tar -C $DATA_DIR/val -czf $DATA_DIR/tao_od_synthetic_val.tar.gz images labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set dataset path <a class=\"anchor\" id=\"head-1.1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_dataset_path = f\"{DATA_DIR}/tao_od_synthetic_train.tar.gz\"\n",
    "eval_dataset_path = f\"{DATA_DIR}/tao_od_synthetic_val.tar.gz\"\n",
    "test_dataset_path = f\"{DATA_DIR}/tao_od_synthetic_val.tar.gz\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload train dataset <a class=\"anchor\" id=\"head-1.2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create train dataset\n",
    "data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "\n",
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"id\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "train_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Update\n",
    "docker_env_vars = {} # Update any variables to be included while triggering Docker run-time like MLOPs variables \n",
    "dataset_information = {\"name\":\"Train dataset\",\n",
    "                       \"description\":\"My train dataset with kitti\",\n",
    "                       \"docker_env_vars\": docker_env_vars}\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/datasets/{train_dataset_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"id\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Upload\n",
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(train_dataset_path)), model_name, \"train\")\n",
    "split_tar_file(train_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}:upload\"\n",
    "\n",
    "    response = requests.post(endpoint, files=files, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"message\" in response.json().keys() and response.json()[\"message\"] == \"Server recieved file and upload process started\"\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload val dataset <a class=\"anchor\" id=\"head-1.3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create eval dataset\n",
    "data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "\n",
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"id\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "eval_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Update\n",
    "docker_env_vars = {} # Update any variables to be included while triggering Docker run-time like MLOPs variables \n",
    "dataset_information = {\"name\":\"Eval dataset\",\n",
    "                       \"description\":\"My eval dataset with kitti\",\n",
    "                       \"docker_env_vars\": docker_env_vars}\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/datasets/{eval_dataset_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"id\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Upload\n",
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(eval_dataset_path)), model_name, \"eval\")\n",
    "split_tar_file(eval_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{eval_dataset_id}:upload\"\n",
    "\n",
    "    response = requests.post(endpoint, files=files, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"message\" in response.json().keys() and response.json()[\"message\"] == \"Server recieved file and upload process started\"\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and upload test dataset <a class=\"anchor\" id=\"head-1.4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create testing dataset for inference\n",
    "ds_format = \"raw\"\n",
    "data = json.dumps({\"type\":ds_type,\"format\":ds_format})\n",
    "\n",
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"id\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "test_dataset_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Upload\n",
    "output_dir = os.path.join(os.path.dirname(os.path.abspath(test_dataset_path)), model_name, \"test\")\n",
    "split_tar_file(test_dataset_path, output_dir)\n",
    "for idx, tar_dataset_path in enumerate(os.listdir(output_dir)):\n",
    "    print(f\"Uploading {idx+1}/{len(os.listdir(output_dir))} tar split\")\n",
    "    files = [(\"file\",open(os.path.join(output_dir, tar_dataset_path),\"rb\"))]\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{test_dataset_id}:upload\"\n",
    "\n",
    "    response = requests.post(endpoint, files=files, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"message\" in response.json().keys() and response.json()[\"message\"] == \"Server recieved file and upload process started\"\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List the created datasets <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "for rsp in response.json():\n",
    "    rsp_keys = rsp.keys()\n",
    "    assert \"id\" in rsp_keys\n",
    "    assert \"type\" in rsp_keys\n",
    "    assert \"format\" in rsp_keys\n",
    "    assert \"name\" in rsp_keys\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose list output\n",
    "print(\"id\\t\\t\\t\\t\\t type\\t\\t\\t format\\t\\t name\")\n",
    "for rsp in response.json():\n",
    "    print(rsp[\"id\"],\"\\t\",rsp[\"type\"],\"\\t\",rsp[\"format\"],\"\\t\\t\",rsp[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Dataset convert Action <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Choose dataset convert action\n",
    "convert_action = \"no_convert\"\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    if model_name in (\"dssd\", \"ssd\", \"retinanet\"):\n",
    "        convert_action = \"convert_and_index\"\n",
    "    elif \"efficientdet\" in model_name:\n",
    "        convert_action = \"convert_\" + model_name\n",
    "    else:\n",
    "        convert_action = \"convert\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}/specs/{convert_action}/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"default\" in response.json().keys()\n",
    "\n",
    "    print(response)\n",
    "    # print(response.json()) ## Uncomment for verbose schema\n",
    "\n",
    "    train_ds_convert_specs = response.json()[\"default\"]\n",
    "\n",
    "    print(json.dumps(train_ds_convert_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    if \"efficientdet\" in model_name:\n",
    "        train_ds_convert_specs[\"dataset_convert\"][\"num_shards\"] = 256\n",
    "        train_ds_convert_specs[\"dataset_convert\"][\"tag\"] = \"train\"\n",
    "    else:\n",
    "        train_ds_convert_specs[\"kitti_config\"][\"image_extension\"] = \".jpg\" #Change to png if your entire dataset is of png format\n",
    "\n",
    "    if convert_action == \"convert_and_index\": # This can be applied for \"convert\" action also for networks like dnv2, yolo's, frcnn etc but not mandatory; for convert_and_index action it is mandatory\n",
    "        # Map your classes to a superclass like mapping pedestrian to person or just the same name\n",
    "        # Mention the classes in the dataset and their mapping\n",
    "        train_ds_convert_specs[\"target_class_mapping\"] = [{\"key\":\"cone\",\"value\":\"cone\"},\n",
    "                                 {\"key\":\"cart\",\"value\":\"cart\"},\n",
    "                                 {\"key\":\"fire_extinguisher\",\"value\":\"fire_extinguisher\"},\n",
    "                                 {\"key\":\"forklift\",\"value\":\"forklift\"}]\n",
    "    print(json.dumps(train_ds_convert_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    parent = None\n",
    "    action = convert_action\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action, \"specs\":train_ds_convert_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    train_ds_convert_id = response.json()\n",
    "    job_map[\"train_dataset_convert_\"+convert_action] = train_ds_convert_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    job_id = train_ds_convert_id\n",
    "    endpoint = f\"{base_url}/datasets/{train_dataset_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True) \n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Eval Dataset convert Action <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Now, repeat the same for the eval dataset\n",
    "# Get default spec schema\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    endpoint = f\"{base_url}/datasets/{eval_dataset_id}/specs/{convert_action}/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"default\" in response.json().keys()\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    eval_ds_convert_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(eval_ds_convert_specs, sort_keys=True, indent=4))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    if \"efficientdet\" in model_name:\n",
    "        eval_ds_convert_specs[\"dataset_convert\"][\"num_shards\"] = 256\n",
    "        eval_ds_convert_specs[\"dataset_convert\"][\"tag\"] = \"val\"\n",
    "    else:\n",
    "        eval_ds_convert_specs[\"kitti_config\"][\"image_extension\"] = \".jpg\" #Change to png if your entire dataset is of png format\n",
    "\n",
    "    if convert_action == \"convert_and_index\": # This can be applied for \"convert\" action also for networks like dnv2, yolo's, frcnn etc but not mandatory; for convert_and_index action it is mandatory\n",
    "        #Map your classes to a superclass like mapping pedestrian to person or just the same name\n",
    "        #Mention the classes in the dataset and their mapping\n",
    "        eval_ds_convert_specs[\"target_class_mapping\"] = [{\"key\":\"cone\",\"value\":\"cone\"},\n",
    "                                 {\"key\":\"cart\",\"value\":\"cart\"},\n",
    "                                 {\"key\":\"fire_extinguisher\",\"value\":\"fire_extinguisher\"},\n",
    "                                 {\"key\":\"forklift\",\"value\":\"forklift\"}]\n",
    "    print(json.dumps(eval_ds_convert_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    parent = job_map[\"train_dataset_convert_\"+convert_action]\n",
    "    action = convert_action\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":eval_ds_convert_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/datasets/{eval_dataset_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    print(response.json())\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    eval_ds_convert_id = response.json()\n",
    "    job_map[\"eval_dataset_convert_\"+convert_action] = eval_ds_convert_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    job_id = eval_ds_convert_id\n",
    "    endpoint = f\"{base_url}/datasets/{eval_dataset_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an experiment <a class=\"anchor\" id=\"head-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "encode_key = \"tlt_encode\"\n",
    "checkpoint_choose_method = \"best_model\"\n",
    "data = json.dumps({\"network_arch\":model_name,\"encryption_key\":encode_key,\"checkpoint_choose_method\":checkpoint_choose_method})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments\"\n",
    "\n",
    "response = requests.post(endpoint,data=data,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"id\" in response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "experiment_id = response.json()[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List experiments <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/experiments\"\n",
    "params = {\"network_arch\": model_name}\n",
    "response = requests.get(endpoint, params=params, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "# print(response.json()) ## Uncomment for verbose list output\n",
    "print(\"model id\\t\\t\\t     network architecture\")\n",
    "for rsp in response.json():\n",
    "    rsp_keys = rsp.keys()\n",
    "    assert \"id\" in rsp_keys and \"network_arch\" in rsp_keys\n",
    "    print(rsp[\"name\"], rsp[\"id\"],rsp[\"network_arch\"])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign train, eval datasets <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "docker_env_vars = {} # Update any variables to be included while triggering Docker run-time like MLOPs variables \n",
    "dataset_information = {\"train_datasets\":[train_dataset_id],\n",
    "                       \"eval_dataset\":eval_dataset_id,\n",
    "                       \"inference_dataset\":test_dataset_id,\n",
    "                       \"calibration_dataset\":train_dataset_id,\n",
    "                       \"docker_env_vars\": docker_env_vars}\n",
    "data = json.dumps(dataset_information)\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Assign PTM <a class=\"anchor\" id=\"head-7\"></a>\n",
    "\n",
    "Search for PTM on NGC for the Object Detection model chosen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List all pretrained models for the chosen network architecture\n",
    "endpoint = f\"{base_url}/experiments\"\n",
    "params = {\"network_arch\": model_name}\n",
    "response = requests.get(endpoint, params=params, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "response_json = response.json()\n",
    "\n",
    "for rsp in response_json:\n",
    "    rsp_keys = rsp.keys()\n",
    "    if \"encryption_key\" not in rsp.keys():\n",
    "        assert \"name\" in rsp_keys and \"version\" in rsp_keys and \"ngc_path\" in rsp_keys and \"additional_id_info\" in rsp_keys\n",
    "        print(f'PTM Name: {rsp[\"name\"]}; PTM version: {rsp[\"version\"]}; NGC PATH: {rsp[\"ngc_path\"]}; Additional info: {rsp[\"additional_id_info\"]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Assigning pretrained models to different object detection models versions\n",
    "# From the output of previous cell make the appropriate changes to this map if you want to change the default PTM backbone.\n",
    "# Changing the default backbone here requires changing default spec/config during train/eval etc like for example\n",
    "# If you are changing the ptm to resnet34, then you have to modify the config key num_layers if it exists to 34 manually\n",
    "pretrained_map = {\"detectnet_v2\" : \"detectnet_v2:resnet18\",\n",
    "                  \"deformable_detr\": \"pretrained_deformable_detr_nvimagenet:resnet50\",\n",
    "                  \"dino\": \"pretrained_dino_nvimagenet:resnet50\",\n",
    "                  \"dssd\" : \"pretrained_object_detection:resnet18\",\n",
    "                  \"efficientdet_tf1\" : \"pretrained_efficientdet:efficientnet_b0\",\n",
    "                  \"efficientdet_tf2\" : \"pretrained_efficientdet_tf2:efficientnet_b0\",\n",
    "                  \"faster_rcnn\" : \"pretrained_object_detection:resnet18\",\n",
    "                  \"retinanet\" : \"pretrained_object_detection:resnet18\",\n",
    "                  \"ssd\" : \"pretrained_object_detection:resnet18\",\n",
    "                  \"yolo_v3\" : \"pretrained_object_detection:resnet18\",\n",
    "                  \"yolo_v4\" : \"pretrained_object_detection:resnet18\",\n",
    "                  \"yolo_v4_tiny\": \"pretrained_object_detection:cspdarknet_tiny\"}\n",
    "no_ptm_models = set([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get pretrained model\n",
    "if model_name not in no_ptm_models:\n",
    "    endpoint = f\"{base_url}/experiments\"\n",
    "    params = {\"network_arch\": model_name}\n",
    "    response = requests.get(endpoint, params=params, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    response_json = response.json()\n",
    "\n",
    "    # Search for ptm with given ngc path\n",
    "    ptm = []\n",
    "    for rsp in response_json:\n",
    "        rsp_keys = rsp.keys()\n",
    "        assert \"ngc_path\" in rsp_keys\n",
    "        if rsp[\"ngc_path\"].endswith(pretrained_map[model_name]):\n",
    "            assert \"id\" in rsp_keys\n",
    "            ptm_id = rsp[\"id\"]\n",
    "            ptm = [ptm_id]\n",
    "            print(\"Metadata for model with requested NGC Path\")\n",
    "            print(rsp)\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name not in no_ptm_models:\n",
    "    ptm_information = {\"base_experiment\":ptm}\n",
    "    data = json.dumps(ptm_information)\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "    response = requests.patch(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "\n",
    "    print(response)\n",
    "    print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View hyperparameters that are enabled for AutoML by default <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Get default spec schema\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/train/schema\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"automl_default_parameters\" in response.json().keys()\n",
    "    automl_specs = response.json()[\"automl_default_parameters\"]\n",
    "    print(json.dumps(automl_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Actions <a class=\"anchor\" id=\"head-10\"></a>\n",
    "\n",
    "For all actions:\n",
    "1. Get default spec schema and derive the default values\n",
    "2. Modify defaults if needed\n",
    "3. Post spec dictionary to the service\n",
    "4. Run model action\n",
    "5. Monitor job using retrieve\n",
    "6. Download results using job download endpoint (if needed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set AutoML related configurations <a class=\"anchor\" id=\"head-9\"></a>\n",
    "Refer to these hyper-links to see the parameters supported by each network and add more parameters if necessary in addition to the default automl enabled parameters:\n",
    "\n",
    "[DetectNet_V2](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/detectnet_v2/detectnet_v2%20-%20train.csv), \n",
    "[Deformable Detr](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/deformable_detr/deformable_detr%20-%20train.csv), \n",
    "[DINO](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/dino/dino%20-%20train.csv), \n",
    "[DSSD](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/dssd/dsssd%20-%20train.csv), \n",
    "[EfficientDet TF1](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/efficientdet_tf1/efficientdet_tf1%20-%20train.csv), \n",
    "[EfficientDet TF2](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/efficientdet_tf2/efficientdet_tf2%20-%20train.csv), \n",
    "[FasterRCNN](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/faster_rcnn/faster_rcnn%20-%20train.csv), \n",
    "[RetinaNet](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/retinanet/retinanet%20-%20train.csv), \n",
    "[SSD](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/ssd/ssd%20-%20train.csv), \n",
    "[YOLO_V3](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/yolo_v3/yolo_v3%20-%20train.csv), \n",
    "[YOLO_V4](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/yolo_v4/yolo_v4%20-%20train.csv), \n",
    "[YOLO_V4_Tiny](https://github.com/NVIDIA/tao_front_end_services/tree/main/api/specs_utils/specs/yolo_v4_tiny/yolo_v4_tiny%20-%20train.csv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if automl_enabled:\n",
    "    # Choose any metric that is present in the kpi dictionary present in the model's status.json. \n",
    "    # Example status.json for each model can be found in the respective section in NVIDIA TAO DOCS here: https://docs.nvidia.com/tao/tao-toolkit/text/model_zoo/cv_models/index.html\n",
    "    metric = \"kpi\"\n",
    "\n",
    "    additional_automl_parameters = [] #Refer to parameter list mentioned in the above links and add any extra parameter in addition to the default enabled ones\n",
    "    remove_default_automl_parameters = [] #Remove any hyperparameters that are enabled by default for AutoML\n",
    "\n",
    "    automl_information = {\"automl_enabled\": True,\n",
    "                          \"automl_algorithm\": automl_algorithm,\n",
    "                          \"metric\": metric,\n",
    "                          \"automl_max_recommendations\": 20, # Only for bayesian\n",
    "                          \"automl_R\": 27, # Only for hyperband\n",
    "                          \"automl_nu\": 3, # Only for hyperband\n",
    "                          \"epoch_multiplier\": 1, # Only for hyperband\n",
    "                          # Enable this if you want to add parameters to automl_add_hyperparameters below that are disabled by TAO in the automl_enabled column of the spec csv.\n",
    "                          # Warning: The parameters that are disabled are not tested by TAO, so there might be unexpected behaviour in overriding this\n",
    "                          \"override_automl_disabled_params\": False,\n",
    "                          \"automl_add_hyperparameters\": str(additional_automl_parameters),\n",
    "                          \"automl_remove_hyperparameters\": str(remove_default_automl_parameters)}\n",
    "    data = json.dumps(automl_information)\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "    response = requests.patch(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    \n",
    "    print(response)\n",
    "    print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/train/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"default\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "train_specs = response.json()[\"default\"]\n",
    "print(json.dumps(train_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if model_name in (\"deformable_detr\", \"dino\"):\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 10\n",
    "    train_specs[\"dataset\"][\"num_classes\"] = int(num_classes) + 1\n",
    "    train_specs[\"train\"][\"num_gpus\"] = 1\n",
    "elif model_name == \"efficientdet_tf1\":\n",
    "    train_specs[\"training_config\"][\"num_epochs\"] = 10\n",
    "    train_specs[\"training_config\"][\"train_batch_size\"] = 2\n",
    "    train_specs[\"training_config\"][\"num_examples_per_epoch\"] = 1414  # number of images in your dataset/number of gpu's\n",
    "    train_specs[\"dataset_config\"][\"num_classes\"] = int(num_classes)  # num_classes was computed during kitti_to_coco_conversion\n",
    "    train_specs[\"eval_config\"][\"eval_epoch_cycle\"] = 2 # Set to number of epochs or less\n",
    "    train_specs[\"gpus\"] = 1\n",
    "elif model_name == \"efficientdet_tf2\":\n",
    "    train_specs[\"train\"][\"num_epochs\"] = 10\n",
    "    train_specs[\"gpus\"] = 1\n",
    "    train_specs[\"train\"][\"lr_schedule\"][\"warmup_epoch\"] = 2 # Set to 1 less than number of epochs or further less\n",
    "    train_specs[\"train\"][\"batch_size\"] = 4\n",
    "    train_specs[\"train\"][\"num_examples_per_epoch\"] = 1414  # number of images in your dataset/number of gpu's\n",
    "    train_specs[\"dataset\"][\"num_classes\"] = int(num_classes)  # num_classes was computed during kitti_to_coco_conversion\n",
    "    train_specs[\"train\"][\"checkpoint_interval\"] = 5 # Set to number of epochs or less\n",
    "else:\n",
    "    # Example for detectnet_v2 (for each network the parameter key might be different)\n",
    "    train_specs[\"training_config\"][\"num_epochs\"] = 10  # num_epochs is the parameter name for all object detection networks\n",
    "    train_specs[\"gpus\"] = 1\n",
    "if \"dataset_config\" in train_specs.keys() and \"image_extension\" in train_specs[\"dataset_config\"].keys():\n",
    "    train_specs[\"dataset_config\"][\"image_extension\"] = \"jpg\"\n",
    "print(json.dumps(train_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map.get(\"eval_dataset_convert_\"+convert_action, job_map.get(\"train_dataset_convert_\"+convert_action, None))\n",
    "parent_id = eval_dataset_id\n",
    "action = \"train\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":train_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert response.json()\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"train_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "# For automl: Training times for different models benchmarked on 1 GPU V100 machine can be found here: https://docs.nvidia.com/tao/tao-toolkit/text/automl/automl.html#results-of-automl-experiments\n",
    "\n",
    "job_id = job_map[\"train_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    if \"error_desc\" in response.json().keys() and response.json()[\"error_desc\"] in (\"Job trying to retrieve not found\", \"No AutoML run found\"):\n",
    "        print(\"Job is being created\")\n",
    "        time.sleep(5)\n",
    "        continue\n",
    "    print(response)\n",
    "    print(json.dumps(response.json(), sort_keys=True, indent=4))\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## To Stop an AutoML JOB\n",
    "#    1. Stop the 'Monitor job status by repeatedly running this cell' cell (the cell right before this cell) manually\n",
    "#    2. Uncomment the snippet in the next cell and run the cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# if automl_enabled:\n",
    "#     job_id = job_map[\"train_\" + model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:cancel\"\n",
    "\n",
    "#     response = requests.post(endpoint, headers=headers)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Resume AutoML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Uncomment the below snippet if you want to resume an already stopped AutoML job and then run the 'Monitor job status by repeatedly running this cell' cell above (4th cell above from this cell)\n",
    "# if automl_enabled:\n",
    "#     job_id = job_map[\"train_\" + model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:resume\"\n",
    "\n",
    "#     data = json.dumps({\"parent_job_id\":parent,\"specs\":train_specs})\n",
    "#     response = requests.post(endpoint, data=data, headers=headers)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download train job artifacts <a class=\"anchor\" id=\"head-12\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Example to list the files of the executed train job\n",
    "job_id = job_map[\"train_\" + model_name]\n",
    "endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:list_files'\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Patch the model with proper metric before training to run this cell; By default loss is used, but some models dont log the parameter under the name 'loss'\n",
    "\n",
    "# # Download selective job contents once the above job shows \"Done\" status\n",
    "# # Example to download selective files of train job (Note: will take time)\n",
    "# endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download_selective_files'\n",
    "\n",
    "# file_lists = [] # Choose file names from the previous cell where all the files for this job were listed\n",
    "# best_model = False # Enable this to download the checkpoint of the best performing model w.r.t to the metric chosen before starting training\n",
    "# latest_model = True # Enable this to download the latest checkpoint of the training job; Disable best_model to use latest_model\n",
    "\n",
    "# params = {\"file_lists\": file_lists, \"best_model\": best_model, \"latest_model\": latest_model}\n",
    "\n",
    "# # Save\n",
    "# temptar = f'{job_id}.tar.gz'\n",
    "# with requests.get(endpoint, headers=headers, params=params, stream=True) as r:\n",
    "#     r.raise_for_status()\n",
    "#     with open(temptar, 'wb') as f:\n",
    "#         for chunk in r.iter_content(chunk_size=8192):\n",
    "#             f.write(chunk)\n",
    "\n",
    "# print(\"Untarring\")\n",
    "# # Untar to destination\n",
    "# tar_command = f'tar -xvf {temptar} -C {workdir}/'\n",
    "# os.system(tar_command)\n",
    "# os.remove(temptar)\n",
    "# print(f\"Results at {workdir}/{job_id}\")\n",
    "# model_downloaded_path = f\"{workdir}/{job_id}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Downloading train job takes a longer time, uncomment this cell if you want to still proceed\n",
    "if download_jobs:\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "    print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "    !python3 -m pip install tqdm\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "    temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "    with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "        while True:\n",
    "            # Check if the file already exists\n",
    "            headers_download_job = dict(headers)\n",
    "            if os.path.exists(temptar):\n",
    "                # Get the current file size\n",
    "                file_size = os.path.getsize(temptar)\n",
    "                print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                # If the file size matches the expected size, break out of the loop\n",
    "                if file_size >= (expected_file_size-1):\n",
    "                    print(\"Download completed successfully.\")\n",
    "                    print(\"Untarring\")\n",
    "                    # Untar to destination\n",
    "                    tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                    os.system(tar_command)\n",
    "                    os.remove(temptar)\n",
    "                    print(f\"Results at {workdir}/{job_id}\")\n",
    "                    model_downloaded_path = f\"{workdir}/{job_id}\"\n",
    "                    break\n",
    "\n",
    "                # Set the headers to resume the download from where it left off\n",
    "                headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "            # Open the file for writing in binary mode\n",
    "            with open(temptar, 'ab') as f:\n",
    "                try:\n",
    "                    response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                    print(response)\n",
    "                    # Check if the request was successful\n",
    "                    if response.status_code in [200, 206]:\n",
    "                        # Iterate over the content in chunks\n",
    "                        for chunk in response.iter_content(chunk_size=1024):\n",
    "                            if chunk:\n",
    "                                # Write the chunk to the file\n",
    "                                f.write(chunk)\n",
    "                                # Flush and sync the file to disk\n",
    "                                f.flush()\n",
    "                                os.fsync(f.fileno())\n",
    "                            progress_bar.update(len(chunk))\n",
    "                    else:\n",
    "                        print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                except requests.exceptions.RequestException as e:\n",
    "                    print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                    time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                    continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# View the checkpoints generated for the training job and for automl jobs, in addition view: best performing model's config and the results of all automl experiments\n",
    "\n",
    "if download_jobs:\n",
    "    if automl_enabled:\n",
    "        !python3 -m pip install pandas==1.5.1\n",
    "        import pandas as pd\n",
    "        model_downloaded_path = f\"{model_downloaded_path}/best_model\"\n",
    "        assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "\n",
    "    assert os.path.exists(model_downloaded_path)\n",
    "    assert (glob.glob(model_downloaded_path + \"/**/*.tlt\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.hdf5\", recursive=True) + glob.glob(model_downloaded_path + \"/**/*.pth\", recursive=True))\n",
    "\n",
    "    if os.path.exists(model_downloaded_path):\n",
    "        #List the binary model file\n",
    "        print(\"\\nCheckpoints for the training experiment\")\n",
    "        if os.path.exists(model_downloaded_path+\"/train/weights\") and len(os.listdir(model_downloaded_path+\"/train/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/train/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/train/weights\"))\n",
    "        elif os.path.exists(model_downloaded_path+\"/weights\") and len(os.listdir(model_downloaded_path+\"/weights\")) > 0:\n",
    "            print(f\"Folder: {model_downloaded_path}/weights\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path+\"/weights\"))\n",
    "        else:\n",
    "            print(f\"Folder: {model_downloaded_path}\")\n",
    "            print(\"Files:\", os.listdir(model_downloaded_path))\n",
    "\n",
    "        if automl_enabled:\n",
    "            assert glob.glob(f\"{model_downloaded_path}/*.protobuf\") or glob.glob(f\"{model_downloaded_path}/*.yaml\")\n",
    "            experiment_artifacts = json.load(open(f\"{model_downloaded_path}/controller.json\",\"r\"))\n",
    "            data_frame = pd.DataFrame(experiment_artifacts)\n",
    "            # Print experiment id/number and the corresponding result\n",
    "            print(\"\\nResults of all experiments\")\n",
    "            with pd.option_context('display.max_rows', None, 'display.max_columns', None, 'display.max_colwidth', None):\n",
    "                print(data_frame[[\"id\",\"result\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate <a class=\"anchor\" id=\"head-12\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get model handler parameters\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "model_parameters = response.json()\n",
    "update_checkpoint_choosing = {}\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = model_parameters[\"checkpoint_choose_method\"]\n",
    "update_checkpoint_choosing[\"checkpoint_epoch_number\"] = model_parameters[\"checkpoint_epoch_number\"]\n",
    "print(update_checkpoint_choosing)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the method by which checkpoint from the parent action is chosen, when parent action is a train/retrain action.\n",
    "# Example for evaluate action below, can be applied in the same way for other actions too\n",
    "update_checkpoint_choosing[\"checkpoint_choose_method\"] = \"latest_model\" # Choose between best_model/latest_model/from_epoch_number\n",
    "# If from_epoch_number is chosen then assign the epoch number to the dictionary key in the format 'from_epoch_number{train_job_id}'\n",
    "# update_checkpoint_choosing[\"checkpoint_epoch_number\"][\"from_epoch_number_28a2754e-50ef-43a8-9733-98913776dd90\"] = 3\n",
    "data = json.dumps(update_checkpoint_choosing)\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.patch(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(json.dumps(response.json(), sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/evaluate/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"default\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "eval_specs = response.json()[\"default\"]\n",
    "print(json.dumps(eval_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name == \"efficientdet_tf1\":\n",
    "    eval_specs[\"dataset_config\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "elif model_name == \"efficientdet_tf2\":\n",
    "    eval_specs[\"evaluate\"][\"num_samples\"] = 353 #number of images in your dataset\n",
    "    eval_specs[\"dataset\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "elif model_name in (\"deformable_detr\", \"dino\"):\n",
    "    eval_specs[\"dataset\"][\"num_classes\"] = int(num_classes) + 1\n",
    "\n",
    "if \"dataset_config\" in eval_specs.keys() and \"image_extension\" in eval_specs[\"dataset_config\"].keys():\n",
    "    eval_specs[\"dataset_config\"][\"image_extension\"] = \"jpg\"\n",
    "print(json.dumps(eval_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map[\"train_\" + model_name]\n",
    "action = \"evaluate\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":eval_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"evaluate_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "job_id = job_map[\"evaluate_\" + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prune, Retrain and Evaluation <a class=\"anchor\" id=\"head-13\"></a>\n",
    "\n",
    "- We optimize the trained model by pruning and retraining in the following cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prune <a class=\"anchor\" id=\"head-14\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/prune/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"default\" in response.json().keys()\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    prune_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(prune_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "# None for prune\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    print(json.dumps(prune_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run actions\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    parent = job_map[\"train_\" + model_name]\n",
    "    action = \"prune\"\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":prune_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    \n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    job_map[\"prune_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell (prune)\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    job_id = job_map['prune_' + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Retrain <a class=\"anchor\" id=\"head-15\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/retrain/schema\"\n",
    "\n",
    "    response = requests.get(endpoint,headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"default\" in response.json().keys()\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    retrain_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes for any of the parameters listed in the previous cell as required\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    # for efficientdet_tf1\n",
    "    if model_name == \"efficientdet_tf1\":\n",
    "        retrain_specs[\"training_config\"][\"num_epochs\"] = 10\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "        retrain_specs[\"training_config\"][\"train_batch_size\"] = 2\n",
    "        retrain_specs[\"training_config\"][\"num_examples_per_epoch\"] = 1414 #number of images in your dataset/number of gpu's\n",
    "        retrain_specs[\"dataset_config\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "        retrain_specs[\"eval_config\"][\"eval_epoch_cycle\"] = 2\n",
    "\n",
    "    # for efficientdet_tf2\n",
    "    elif model_name == \"efficientdet_tf2\":\n",
    "        retrain_specs[\"train\"][\"num_epochs\"] = 10\n",
    "        retrain_specs[\"train\"][\"lr_schedule\"][\"warmup_epoch\"] = 2\n",
    "        retrain_specs[\"train\"][\"checkpoint_interval\"] = 5\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "        retrain_specs[\"train\"][\"batch_size\"] = 4\n",
    "        retrain_specs[\"train\"][\"num_examples_per_epoch\"] = 1414 #number of images in your dataset/number of gpu's\n",
    "        retrain_specs[\"dataset\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "\n",
    "    # Example for (for each network the parameter key might be different)\n",
    "    else:\n",
    "        retrain_specs[\"training_config\"][\"num_epochs\"] = 10 # num_epochs is the parameter name for all object detection networks\n",
    "        retrain_specs[\"gpus\"] = 1\n",
    "\n",
    "    if \"dataset_config\" in retrain_specs.keys() and \"image_extension\" in retrain_specs[\"dataset_config\"].keys():\n",
    "        retrain_specs[\"dataset_config\"][\"image_extension\"] = \"jpg\"\n",
    "    print(json.dumps(retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run actions\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    parent = job_map[\"prune_\" + model_name]\n",
    "    action = \"retrain\"\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":retrain_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    \n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    job_map[\"retrain_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell (retrain)\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    job_id = job_map['retrain_' + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional cancel job - for jobs that are pending/running (retrain)\n",
    "\n",
    "# if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "#     job_id = job_map[\"retrain_\" + model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}:cancel\"\n",
    "\n",
    "#     response = requests.post(endpoint, headers=headers)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional delete job - for jobs that are error/done (retrain)\n",
    "\n",
    "# if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "#     job_id = job_map[\"retrain_\" + model_name]\n",
    "#     endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "#     response = requests.delete(endpoint, headers=headers)\n",
    "\n",
    "#     print(response)\n",
    "#     print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate after retrain <a class=\"anchor\" id=\"head-15\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/specs/evaluate/schema\"\n",
    "\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"default\" in response.json().keys()\n",
    "\n",
    "    print(response)\n",
    "    #print(response.json()) ## Uncomment for verbose schema\n",
    "    eval_retrain_specs = response.json()[\"default\"]\n",
    "    print(json.dumps(eval_retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    if model_name == \"efficientdet_tf1\":\n",
    "        eval_retrain_specs[\"dataset_config\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "    elif model_name == \"efficientdet_tf2\":\n",
    "        eval_retrain_specs[\"evaluate\"][\"num_samples\"] = 353 #number of images in your dataset\n",
    "        eval_retrain_specs[\"dataset\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "    elif model_name in (\"deformable_detr\", \"dino\"):\n",
    "        eval_retrain_specs[\"dataset\"][\"num_classes\"] = int(num_classes) + 1\n",
    "\n",
    "    if \"dataset_config\" in eval_retrain_specs.keys() and \"image_extension\" in eval_retrain_specs[\"dataset_config\"].keys():\n",
    "        eval_retrain_specs[\"dataset_config\"][\"image_extension\"] = \"jpg\"\n",
    "    print(json.dumps(eval_retrain_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run actions\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    parent = job_map[\"retrain_\" + model_name]\n",
    "    action = \"evaluate\"\n",
    "    data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":eval_retrain_specs})\n",
    "\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "    response = requests.post(endpoint, data=data, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    \n",
    "    print(response)\n",
    "    print(response.json())\n",
    "\n",
    "    job_map[\"eval_retrain_\" + model_name] = response.json()\n",
    "    print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell (evaluate)\n",
    "if model_name not in (\"deformable_detr\", \"dino\"):\n",
    "    job_id = job_map['eval_retrain_' + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "    while True:\n",
    "        clear_output(wait=True)\n",
    "        response = requests.get(endpoint, headers=headers)\n",
    "        assert response.status_code in (200, 201)\n",
    "        assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "        print(response)\n",
    "        print(response.json())\n",
    "        if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "            break\n",
    "        time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Export <a class=\"anchor\" id=\"head-17\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/export/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"default\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "export_specs = response.json()[\"default\"]\n",
    "print(json.dumps(export_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name == \"efficientdet_tf2\":\n",
    "    export_specs[\"dataset\"][\"num_classes\"] = int(num_classes)\n",
    "elif model_name in (\"deformable_detr\", \"dino\"):\n",
    "    export_specs[\"dataset\"][\"num_classes\"] = int(num_classes) + 1\n",
    "print(json.dumps(export_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map[\"train_\" + model_name]\n",
    "action = \"export\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":export_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"export_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "job_id = job_map['export_' + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT Engine generation using TAO-Deploy <a class=\"anchor\" id=\"head-19\"></a>\n",
    "\n",
    "- Here, we use the exported model to generate trt engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/gen_trt_engine/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"default\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "tao_deploy_specs = response.json()[\"default\"]\n",
    "print(json.dumps(tao_deploy_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes\n",
    "if model_name == \"efficientdet_tf2\":\n",
    "    tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"int8\"\n",
    "    tao_deploy_specs[\"dataset\"][\"num_classes\"] = int(num_classes)\n",
    "elif model_name in (\"deformable_detr\", \"dino\"):\n",
    "    tao_deploy_specs[\"gen_trt_engine\"][\"tensorrt\"][\"data_type\"] = \"int8\"\n",
    "    tao_deploy_specs[\"dataset\"][\"num_classes\"] = int(num_classes) + 1\n",
    "elif model_name == \"dssd\":\n",
    "    tao_deploy_specs[\"data_type\"] = \"fp16\"\n",
    "else:\n",
    "    tao_deploy_specs[\"data_type\"] = \"int8\"\n",
    "print(json.dumps(tao_deploy_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map[\"export_\" + model_name]\n",
    "action = \"gen_trt_engine\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":tao_deploy_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"model_gen_trt_engine_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "job_id = job_map['model_gen_trt_engine_' + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TAO inference <a class=\"anchor\" id=\"head-20\"></a>\n",
    "\n",
    "- Run inference on a set of images using the .tlt model created at train step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/inference/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"default\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "tao_inference_specs = response.json()[\"default\"]\n",
    "print(json.dumps(tao_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to the specs dictionary if necessary\n",
    "if model_name == \"efficientdet_tf1\":\n",
    "    tao_inference_specs[\"dataset_config\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "elif model_name == \"efficientdet_tf2\":\n",
    "    tao_inference_specs[\"dataset\"][\"num_classes\"] = int(num_classes)\n",
    "elif model_name in (\"deformable_detr\", \"dino\"):\n",
    "    tao_inference_specs[\"dataset\"][\"num_classes\"] = int(num_classes) + 1\n",
    "\n",
    "if \"dataset_config\" in tao_inference_specs.keys() and \"image_extension\" in tao_inference_specs[\"dataset_config\"].keys():\n",
    "    tao_inference_specs[\"dataset_config\"][\"image_extension\"] = \"jpg\"\n",
    "print(json.dumps(tao_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map[\"train_\" + model_name]\n",
    "action = \"inference\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":tao_inference_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"inference_tlt_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "job_id = job_map['inference_tlt_' + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Download job contents once the above job shows \"Done\" status\n",
    "if download_jobs:\n",
    "    job_id = job_map[\"inference_tlt_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "    print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "    !python3 -m pip install tqdm\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "    temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "    with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "        while True:\n",
    "            # Check if the file already exists\n",
    "            headers_download_job = dict(headers)\n",
    "            if os.path.exists(temptar):\n",
    "                # Get the current file size\n",
    "                file_size = os.path.getsize(temptar)\n",
    "                print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                # If the file size matches the expected size, break out of the loop\n",
    "                if file_size >= (expected_file_size-1):\n",
    "                    print(\"Download completed successfully.\")\n",
    "                    print(\"Untarring\")\n",
    "                    # Untar to destination\n",
    "                    tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                    os.system(tar_command)\n",
    "                    os.remove(temptar)\n",
    "                    print(f\"Results at {workdir}/{job_id}\")\n",
    "                    inference_out_path = f\"{workdir}/{job_id}\"\n",
    "                    break\n",
    "\n",
    "                # Set the headers to resume the download from where it left off\n",
    "                headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "            # Open the file for writing in binary mode\n",
    "            with open(temptar, 'ab') as f:\n",
    "                try:\n",
    "                    response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                    print(response)\n",
    "                    # Check if the request was successful\n",
    "                    if response.status_code in [200, 206]:\n",
    "                        # Iterate over the content in chunks\n",
    "                        for chunk in response.iter_content(chunk_size=1024):\n",
    "                            if chunk:\n",
    "                                # Write the chunk to the file\n",
    "                                f.write(chunk)\n",
    "                                # Flush and sync the file to disk\n",
    "                                f.flush()\n",
    "                                os.fsync(f.fileno())\n",
    "                            progress_bar.update(len(chunk))\n",
    "                    else:\n",
    "                        print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                except requests.exceptions.RequestException as e:\n",
    "                    print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                    time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                    continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    from IPython.display import Image\n",
    "    assert glob.glob(f\"{inference_out_path}/**/*.jpg\", recursive=True)\n",
    "    sample_image = glob.glob(f\"{inference_out_path}/**/*.jpg\", recursive=True)[6]\n",
    "    Image(filename=sample_image) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TRT inference <a class=\"anchor\" id=\"head-21\"></a>\n",
    "\n",
    "- no need to change the specs since we already uploaded it at the tlt inference step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get default spec schema\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/specs/inference/schema\"\n",
    "\n",
    "response = requests.get(endpoint, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "assert \"default\" in response.json().keys()\n",
    "\n",
    "print(response)\n",
    "#print(response.json()) ## Uncomment for verbose schema\n",
    "trt_inference_specs = response.json()[\"default\"]\n",
    "print(json.dumps(trt_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Apply changes to the specs dictionary if necessary\n",
    "if model_name == \"efficientdet_tf1\":\n",
    "    trt_inference_specs[\"dataset_config\"][\"num_classes\"] = int(num_classes) #num_classes was computed during kitti_to_coco_conversion\n",
    "elif model_name == \"efficientdet_tf2\":\n",
    "    trt_inference_specs[\"dataset\"][\"num_classes\"] = int(num_classes)\n",
    "elif model_name in (\"deformable_detr\", \"dino\"):\n",
    "    trt_inference_specs[\"dataset\"][\"num_classes\"] = int(num_classes) + 1\n",
    "\n",
    "if \"dataset_config\" in trt_inference_specs.keys() and \"image_extension\" in trt_inference_specs[\"dataset_config\"].keys():\n",
    "    trt_inference_specs[\"dataset_config\"][\"image_extension\"] = \"jpg\"\n",
    "print(json.dumps(trt_inference_specs, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run action\n",
    "parent = job_map[\"model_gen_trt_engine_\" + model_name]\n",
    "action = \"inference\"\n",
    "data = json.dumps({\"parent_job_id\":parent,\"action\":action,\"specs\":trt_inference_specs})\n",
    "\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs\"\n",
    "\n",
    "response = requests.post(endpoint, data=data, headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())\n",
    "\n",
    "job_map[\"inference_trt_\" + model_name] = response.json()\n",
    "print(job_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Monitor job status by repeatedly running this cell\n",
    "job_id = job_map['inference_trt_' + model_name]\n",
    "endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "\n",
    "while True:\n",
    "    clear_output(wait=True)\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    assert \"status\" in response.json().keys() and response.json().get(\"status\") != \"Error\"\n",
    "    print(response)\n",
    "    print(response.json())\n",
    "    if response.json().get(\"status\") in [\"Done\",\"Error\", \"Canceled\"] or response.status_code not in (200,201):\n",
    "        break\n",
    "    time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Download job contents once the above job shows \"Done\" status\n",
    "if download_jobs:\n",
    "    job_id = job_map[\"inference_trt_\" + model_name]\n",
    "    endpoint = f\"{base_url}/experiments/{experiment_id}/jobs/{job_id}\"\n",
    "    response = requests.get(endpoint, headers=headers)\n",
    "    assert response.status_code in (200, 201)\n",
    "    expected_file_size = response.json().get(\"job_tar_stats\", {}).get(\"file_size\")\n",
    "    print(\"expected_file_size: \", expected_file_size)\n",
    "\n",
    "    !python3 -m pip install tqdm\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    endpoint = f'{base_url}/experiments/{experiment_id}/jobs/{job_id}:download'\n",
    "    temptar = f'{job_id}.tar.gz'\n",
    "\n",
    "    with tqdm(total=expected_file_size, unit='B', unit_scale=True) as progress_bar:\n",
    "        while True:\n",
    "            # Check if the file already exists\n",
    "            headers_download_job = dict(headers)\n",
    "            if os.path.exists(temptar):\n",
    "                # Get the current file size\n",
    "                file_size = os.path.getsize(temptar)\n",
    "                print(f\"File size of dowloaded content until now is {file_size}\")\n",
    "\n",
    "                # If the file size matches the expected size, break out of the loop\n",
    "                if file_size >= (expected_file_size-1):\n",
    "                    print(\"Download completed successfully.\")\n",
    "                    print(\"Untarring\")\n",
    "                    # Untar to destination\n",
    "                    tar_command = f'tar -xf {temptar} -C {workdir}/'\n",
    "                    os.system(tar_command)\n",
    "                    os.remove(temptar)\n",
    "                    print(f\"Results at {workdir}/{job_id}\")\n",
    "                    inference_out_path = f\"{workdir}/{job_id}\"\n",
    "                    break\n",
    "\n",
    "                # Set the headers to resume the download from where it left off\n",
    "                headers_download_job['Range'] = f'bytes={file_size}-'\n",
    "            # Open the file for writing in binary mode\n",
    "            with open(temptar, 'ab') as f:\n",
    "                try:\n",
    "                    response = requests.get(endpoint, headers=headers_download_job, stream=True)\n",
    "                    print(response)\n",
    "                    # Check if the request was successful\n",
    "                    if response.status_code in [200, 206]:\n",
    "                        # Iterate over the content in chunks\n",
    "                        for chunk in response.iter_content(chunk_size=1024):\n",
    "                            if chunk:\n",
    "                                # Write the chunk to the file\n",
    "                                f.write(chunk)\n",
    "                                # Flush and sync the file to disk\n",
    "                                f.flush()\n",
    "                                os.fsync(f.fileno())\n",
    "                            progress_bar.update(len(chunk))\n",
    "                    else:\n",
    "                        print(f\"Failed to download file. Status code: {response.status_code}\")\n",
    "                except requests.exceptions.RequestException as e:\n",
    "                    print(\"Connection interrupted during download, resuming download from breaking point\")\n",
    "                    time.sleep(5)  # Sleep for a while before retrying the request\n",
    "                    continue  # Continue the loop to retry the request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if download_jobs:\n",
    "    from IPython.display import Image\n",
    "    assert glob.glob(f\"{inference_out_path}/**/*.jpg\", recursive=True)\n",
    "    sample_image = glob.glob(f\"{inference_out_path}/**/*.jpg\", recursive=True)[6]\n",
    "    Image(filename=sample_image) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete model <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/experiments/{experiment_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete train dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets/{train_dataset_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete val dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets/{eval_dataset_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete test dataset <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = f\"{base_url}/datasets/{test_dataset_id}\"\n",
    "\n",
    "response = requests.delete(endpoint,headers=headers)\n",
    "assert response.status_code in (200, 201)\n",
    "\n",
    "print(response)\n",
    "print(response.json())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
